package com.chatplus.application.common.util;

import com.knuddels.jtokkit.Encodings;
import com.knuddels.jtokkit.api.Encoding;
import com.knuddels.jtokkit.api.EncodingRegistry;
import com.knuddels.jtokkit.api.EncodingType;

import java.util.Optional;

public class MessageTokenUtil {
    private MessageTokenUtil() {
    }
    private static final EncodingRegistry registry = Encodings.newDefaultEncodingRegistry();
    public static Long countMessageTokens(String text, String model) {
        Optional<Encoding> optionalEncoding = registry.getEncodingForModel(model);
        Encoding encoding = optionalEncoding.orElse(registry.getEncoding(EncodingType.CL100K_BASE));
        return (long) encoding.countTokensOrdinary(text);
    }
}
